package llm

import (
	"crypto/hmac"
	"crypto/sha256"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"net/url"
	"time"

	"github.com/gorilla/websocket"
)

/* 讯飞星火 */

const (
	xf_app_id  = ""
	xf_api_key = ""

	xf_url = "wss://spark-api.xf-yun.com/v3.5/chat"
)

var xf_api_secret = []byte("")

type Spark struct {
	// The request URL consists of current date, so it should be re-generated for every request.
	url string
}

func (s *Spark) RefreshToken() error {
	// no token at all, generate request url
	s.url = generateRequestUrl()
	return nil
}

func (s *Spark) CallApi(system_prompt, user_prompt string) (string, error) {

	conn, rsp, err := websocket.DefaultDialer.Dial(s.url, nil)
	if err != nil {
		return "", err
	}

	if rsp.StatusCode != 101 {
		return "", fmt.Errorf("websocket connect error, status code: %d", rsp.StatusCode)
	}

	defer conn.Close()

	err = conn.WriteJSON(generateSparkReq(system_prompt, user_prompt))
	if err != nil {
		return "", err
	}

	result := ""
	for {
		_, message, err := conn.ReadMessage()

		if err != nil {
			return "", err
		}

		var rsp SparkRsp
		err = json.Unmarshal(message, &rsp)
		if err != nil {
			return "", err
		}

		if rsp.Header.Code != 0 {
			return "", fmt.Errorf("spark api error, code: %d, message: %s", rsp.Header.Code, rsp.Header.Message)
		}

		result += rsp.Payload.Choices.Text[0].Content

		if rsp.Header.Status == 2 {
			break
		}
	}

	return result, nil
}

func generateRequestUrl() string {

	date := time.Now().UTC().Format(time.RFC1123)

	u, _ := url.Parse(xf_url)

	sign := sign(fmt.Sprintf("host: %s\ndate: %s\nGET %s HTTP/1.1", u.Host, date, u.Path))

	authSrc := fmt.Sprintf(
		"hmac username=\"%s\", algorithm=\"hmac-sha256\", headers=\"host date request-line\", signature=\"%s\"",
		xf_api_key,
		sign)
	authorization := base64.StdEncoding.EncodeToString([]byte(authSrc))

	v := url.Values{}
	v.Add("host", u.Host)
	v.Add("date", date)
	v.Add("authorization", authorization)

	return xf_url + "?" + v.Encode()
}

func sign(src string) string {
	mac := hmac.New(sha256.New, xf_api_secret)
	mac.Write([]byte(src))
	return base64.StdEncoding.EncodeToString(mac.Sum(nil))
}

func generateSparkReq(system_prompt, user_prompt string) map[string]interface{} {

	return map[string]interface{}{

		"header": map[string]string{
			"app_id": xf_app_id,
		},

		"parameter": map[string]interface{}{
			"chat": map[string]string{
				"domain": "generalv3.5",
			},
		},

		"payload": map[string]interface{}{
			"message": map[string]interface{}{
				"text": []map[string]string{
					{"role": "system", "content": system_prompt},
					{"role": "user", "content": user_prompt},
				},
			},
		},
	}
}

type SparkRsp struct {
	Header  SparkHeader  `json:"header"`
	Payload SparkPayload `json:"payload"`
}

type SparkHeader struct {
	Code    int    `json:"code"`
	Message string `json:"message"`
	Sid     string `json:"sid"`
	Status  int    `json:"status"`
}

type SparkPayload struct {
	Choices SparkChoice `json:"choices"`
	Usage   *SparkUsage `json:"usage,omitempty"`
}

type SparkChoice struct {
	Status int         `json:"status"`
	Seq    int         `json:"seq"`
	Text   []SparkText `json:"text"`
}

type SparkText struct {
	Content string `json:"content"`
	Role    string `json:"role"`
	Index   int    `json:"index"`
}

type SparkUsage struct {
	QuestionTokens   int `json:"question_tokens"`
	PromptTokens     int `json:"prompt_tokens"`
	CompletionTokens int `json:"completion_tokens"`
	TotalTokens      int `json:"total_tokens"`
}
